// Date:   Sun Aug 31 03:19:39 PM 2025
// Mail:   lunar_ubuntu@qq.com
// Author: https://github.com/xiaoqixian

#include "av.hpp"
#include <stdexcept>
#include <opencv2/opencv.hpp>
extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/opt.h>
#include <libavutil/time.h>
#include <libswscale/swscale.h>
}

AVContext::AVContext(std::string const& f): filepath_(f) {
  if (avformat_open_input(&fmt_ctx_, filepath_.c_str(), nullptr, nullptr) < 0) {
    throw std::runtime_error("Failed to open input video");
  }
  if (avformat_find_stream_info(fmt_ctx_, nullptr) < 0) {
    throw std::runtime_error("Failed to find stream info");
  }
  // Find the first video stream
  for (size_t i = 0; i < fmt_ctx_->nb_streams; i++) {
    if (fmt_ctx_->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
      video_stream_index_ = i;
      break;
    }
  }
  if (video_stream_index_ < 0) throw std::runtime_error("No video stream found");

  // Initialize decoder
  codec_par_ = fmt_ctx_->streams[video_stream_index_]->codecpar;
  codec_ = avcodec_find_decoder(codec_par_->codec_id);
  if (codec_ == nullptr) {
    throw std::runtime_error("Failed to find a decoder");
  }
}

AVContext::~AVContext() {
  if (sws_ctx_) {
    sws_freeContext(sws_ctx_);
    sws_ctx_ = nullptr;
  }
  if (fmt_ctx_) {
    avformat_close_input(&fmt_ctx_);
  }
}

AVContext::AVContext(AVContext&& other) noexcept
  : filepath_(std::move(other.filepath_)),
    fmt_ctx_(std::exchange(other.fmt_ctx_, nullptr)),
    codec_par_(std::exchange(other.codec_par_, nullptr)),
    codec_(std::exchange(other.codec_, nullptr)),
    codec_ctx_map_(std::move(other.codec_ctx_map_)),
    video_stream_index_(std::exchange(other.video_stream_index_, -1))
{}

AVContext& AVContext::operator=(AVContext&& other) noexcept {
  if (this != &other) {
    if (fmt_ctx_) avformat_close_input(&fmt_ctx_);

    filepath_ = std::move(other.filepath_);
    fmt_ctx_ = std::exchange(other.fmt_ctx_, nullptr);
    codec_ctx_map_ = std::move(other.codec_ctx_map_);
    codec_par_ = std::exchange(other.codec_par_, nullptr);
    codec_ = std::exchange(other.codec_, nullptr);
    video_stream_index_ = std::exchange(other.video_stream_index_, -1);
  }
  return *this;
}

void AVContext::rewind() {
  av_seek_frame(fmt_ctx_, video_stream_index_, 0, AVSEEK_FLAG_BACKWARD);
}

std::vector<int64_t> compile_videos(const std::vector<std::string>& video_paths, const char* out_path) {
  const int width = 1280;
  const int height = 720;
  // const double fps = 30.0;
  const int fps = 30;
  std::vector<int64_t> segments;
  segments.reserve(video_paths.size());
  std::vector<size_t> seg_frame_cnt;
  seg_frame_cnt.reserve(video_paths.size()+1);

  {
    av_log_set_level(AV_LOG_ERROR);
    avformat_network_init();
    AVFormatContext* fmt_ctx = nullptr;
    avformat_alloc_output_context2(&fmt_ctx, nullptr, nullptr, out_path);
    if (!fmt_ctx) {
      throw std::runtime_error("avformat_alloc_output_context2 error\n");
    }

    AVCodec* codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!codec) throw std::runtime_error("H.264 codec not found");

    AVStream* stream = avformat_new_stream(fmt_ctx, codec);
    if (!stream) throw std::runtime_error("Failed to create stream");

    AVCodecContext* codec_ctx = avcodec_alloc_context3(codec);
    codec_ctx->codec_id = AV_CODEC_ID_H264;
    codec_ctx->bit_rate = 4000000;
    codec_ctx->width = width;
    codec_ctx->height = height;
    codec_ctx->time_base = AVRational{1, fps};
    codec_ctx->framerate = AVRational{fps, 1};
    codec_ctx->gop_size = fps / 2; // max keyframe interval
    codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;

    stream->time_base = codec_ctx->time_base;

    if (fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
        codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

    if (avcodec_open2(codec_ctx, codec, nullptr) < 0)
        throw std::runtime_error("Failed to open codec");

    avcodec_parameters_from_context(stream->codecpar, codec_ctx);

    if (!(fmt_ctx->oformat->flags & AVFMT_NOFILE)) {
        if (avio_open(&fmt_ctx->pb, out_path, AVIO_FLAG_WRITE) < 0)
            throw std::runtime_error("Failed to open output file");
    }

    int ret = avformat_write_header(fmt_ctx, nullptr);
    if (ret < 0) throw std::runtime_error("");

    SwsContext* sws_ctx = sws_getContext(
        width, height, AV_PIX_FMT_BGR24,
        width, height, AV_PIX_FMT_YUV420P,
        SWS_BILINEAR, nullptr, nullptr, nullptr);

    AVFrame* av_frame = av_frame_alloc();
    av_frame->format = AV_PIX_FMT_YUV420P;
    av_frame->width = width;
    av_frame->height = height;
    av_frame_get_buffer(av_frame, 0);

    int64_t pts = 0;
    size_t frame_cnt = 0;
    // bool new_segment;
    for (size_t i = 0; i < video_paths.size(); ++i) {
      // new_segment = true;
      if (i > 0) {
        seg_frame_cnt.emplace_back(frame_cnt);
      }

      cv::VideoCapture cap(video_paths[i]);
      if (!cap.isOpened())
        throw std::runtime_error("Failed to open input video: " + video_paths[i]);

      cv::Mat cv_frame, cv_resized;
      while (true) {
        if (!cap.read(cv_frame)) break;
        frame_cnt++;

        cv::resize(cv_frame, cv_resized, cv::Size(width, height));

        // convert BGR -> YUV420P
        const uint8_t* in_data[1] = {cv_resized.data};
        int in_linesize[1] = {static_cast<int>(cv_resized.step[0])};
        sws_scale(sws_ctx, in_data, in_linesize, 0, height, av_frame->data, av_frame->linesize);

        av_frame->pts = av_rescale_q(pts, {1, fps}, codec_ctx->time_base);
        pts++;
        AVPacket pkt;
        av_init_packet(&pkt);
        pkt.data = nullptr;
        pkt.size = 0;

        if (avcodec_send_frame(codec_ctx, av_frame) < 0) {
          throw std::runtime_error("error sending frame");
        }

        while (avcodec_receive_packet(codec_ctx, &pkt) == 0) {
          av_packet_rescale_ts(&pkt, codec_ctx->time_base, stream->time_base);
          pkt.stream_index = stream->index;
          // if (new_segment) {
          //   segments.emplace_back(pkt.pts);
          //   new_segment = false;
          // }
          av_interleaved_write_frame(fmt_ctx, &pkt);
          av_packet_unref(&pkt);
        }
      }

      cap.release();
    }
    printf("all videos done\n");

    // flush encoder
    avcodec_send_frame(codec_ctx, nullptr);
    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.data = nullptr;
    pkt.size = 0;
    while (avcodec_receive_packet(codec_ctx, &pkt) == 0) {
      av_packet_rescale_ts(&pkt, codec_ctx->time_base, stream->time_base);
      pkt.stream_index = stream->index;
      av_interleaved_write_frame(fmt_ctx, &pkt);
      av_packet_unref(&pkt);
    }
    av_write_trailer(fmt_ctx);
    av_frame_free(&av_frame);
    sws_freeContext(sws_ctx);
    avcodec_free_context(&codec_ctx);
    if (!(fmt_ctx->oformat->flags & AVFMT_NOFILE)) avio_closep(&fmt_ctx->pb);
    avformat_free_context(fmt_ctx);
  }

  {
    AVFormatContext* fmt_ctx = nullptr;
    if (avformat_open_input(&fmt_ctx, out_path, nullptr, nullptr) < 0) {
      throw std::runtime_error("Failed to open input file");
    }

    if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {
      throw std::runtime_error("Failed to find stream info");
    }

    int video_stream_index = -1;
    for (unsigned i = 0; i < fmt_ctx->nb_streams; ++i) {
      if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
        video_stream_index = i;
        break;
      }
    }
    if (video_stream_index == -1) {
      throw std::runtime_error("No video stream found");
    }

    AVStream* video_stream = fmt_ctx->streams[video_stream_index];

    AVCodec* codec = avcodec_find_decoder(video_stream->codecpar->codec_id);
    if (!codec) {
      throw std::runtime_error("Codec not found");
    }

    AVCodecContext* codec_ctx = avcodec_alloc_context3(codec);
    if (!codec_ctx) {
      throw std::runtime_error("Failed to allocate codec context");
    }

    if (avcodec_parameters_to_context(codec_ctx, video_stream->codecpar) < 0) {
      throw std::runtime_error("Failed to copy codec parameters");
    }

    if (avcodec_open2(codec_ctx, codec, nullptr) < 0) {
      throw std::runtime_error("Failed to open codec");
    }

    AVPacket* pkt = av_packet_alloc();
    AVFrame* frame = av_frame_alloc();

    using P = int64_t;
    std::vector<P> temp_frame_loc;

    while (av_read_frame(fmt_ctx, pkt) >= 0 && !seg_frame_cnt.empty()) {
      if (pkt->stream_index != video_stream_index) {
        av_packet_unref(pkt);
        continue;
      }

      if (avcodec_send_packet(codec_ctx, pkt) < 0) {
        av_packet_unref(pkt);
        continue;
      }

      while (avcodec_receive_frame(codec_ctx, frame) == 0) {
        temp_frame_loc.emplace_back(frame->pts);
        av_frame_unref(frame);
      }

      av_packet_unref(pkt);
    }

    // flush decoder
    avcodec_send_packet(codec_ctx, nullptr);
    while (!seg_frame_cnt.empty() && avcodec_receive_frame(codec_ctx, frame) == 0) {
      temp_frame_loc.emplace_back(frame->pts);
      av_frame_unref(frame);
    }

    av_frame_free(&frame);
    av_packet_free(&pkt);
    avcodec_free_context(&codec_ctx);
    avformat_close_input(&fmt_ctx); 

    std::sort(temp_frame_loc.begin(), temp_frame_loc.end());
    segments.emplace_back(temp_frame_loc[0]);
    for (auto const frame_idx: seg_frame_cnt) {
      assert(frame_idx < temp_frame_loc.size());
      segments.emplace_back(temp_frame_loc[frame_idx]);
    }
  }
  printf("segments: [");
  for (size_t i = 0; i < segments.size(); i++) {
    printf("%ld ", segments[i]);
  }
  printf("]\n");
  assert(segments.size() == video_paths.size());
  
  return segments;
}
